"""
 注意：：：这个版本 测试如果不使用open ai 的大模型 返回的格式不会遵循json（应该是.with_structured_output(xxx)不支持） ，导致运行出错。所以另外基于这个改了一下提示词版本的。

reAct 大模型的一种思维模式，让大模型先思考再行动。（Reason + Act） 框架。
主要工具：
    tavily 搜索检索 ，需要配置api
    langsmit  agent 监控工具 ，需要配置api

本示例：
    使用 LangGraph 构建一个完整的“ReAct智能体工作流”。
    - Tavily 搜索工具：让模型具备网络搜索能力。
    - LangSmith（可选）：用于可视化 agent 的执行过程。

"""
import asyncio
import operator
from typing import TypedDict, List, Annotated, Tuple, Union, Literal

from dotenv import load_dotenv
from langchain import hub
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent
from pydantic import Field, BaseModel

# 自动加载 .env 文件中的环境变量
load_dotenv()

# 创建 Tavily 搜索工具（可联网搜索），每次只返回一个搜索结果
tools = [TavilySearchResults(max_results=1)]

# 从 LangChain Hub 获取 ReAct 框架提示模板
# Hub 上保存了一些通用的 prompt，可以直接 pull 使用
prompt = hub.pull("wfh/react-agent-executor")
prompt.pretty_print()

# ----------------------------------------------------------
# 二、模型定义
# ----------------------------------------------------------
#  讯飞星火 "xdeepseekv32exp"
model_name = "xdeepseekv32exp"

# 定义大语言模型
llm = ChatOpenAI(model=model_name)

# 创建 ReAct 智能体执行器
# 会自动管理“思考 → 行动 → 观察 → 再思考”的循环逻辑
agent_executor = create_react_agent(llm, tools, prompt=prompt)


# ----------------------------------------------------------
# 三、定义状态结构（State）
# ----------------------------------------------------------
# LangGraph 中每个节点执行后会更新 State（类似状态机的上下文）
# PlanExecute 用于在节点间传递任务信息
class PlanExecute(TypedDict):
    input: str  # 用户录入提示词
    plan: List[str]  # 用户录入提示词如果是大问题，会被拆分成几个小问题
    past_steps: Annotated[List[Tuple], operator.add]  # 已执行的步骤及结果
    response: str  # 最终回复给用户的答案


# ----------------------------------------------------------
# 四、定义 Pydantic 数据模型
# ----------------------------------------------------------
# 定义一个计划模型类，用于描述未来要执行的计划
class Plan(BaseModel):
    """表示一组待执行的计划步骤"""
    steps: List[str] = Field(description="需要执行的不同步骤，应该按顺序排列")


# 定义生成计划的提示模板（系统角色）
planner_prompt = ChatPromptTemplate.from_messages([
    ("system",
     """
       对于给定的目标，提出一个简单的逐步计划。
       - 每个步骤应是独立、可执行的任务；
       - 不要添加多余步骤；
       - 最后一步应能直接得到最终答案；
       - 确保每步信息足够，不要跳过。
       """),
    ("placeholder", "{messages}")  # 这里插入用户输入
])

# 构建 planner（计划生成器）
# planner = planner_prompt | ChatOpenAI(model=model_name, temperature=0).with_structured_output(Plan)
planner = planner_prompt | ChatOpenAI(model=model_name, temperature=0).with_structured_output(Plan)

planner.invoke({"input": "2024年奥运会金牌获得最多的人是哪个国家的？"})
